import numpy as np
import matplotlib.pyplot as plt

class CalculateAverages:
    def __init__(self, filenames):
        self.filenames = filenames
        self.num_slots = 10000  # 一共有一万个时隙

    def calculate_slot_averages(self):
        # 初始化列表用于存储每个时隙的数据
        slot_data = [[] for _ in range(self.num_slots)]

        # 读取每个文件中的数据并存储到对应时隙的列表中
        # 读取每个文件中的数据并存储到对应时隙的列表中
        for filename in self.filenames:
            with open(filename, 'r') as file:
                lines = file.readlines()
                for idx, line in enumerate(lines):
                    # 将每行数据转换为浮点数
                    value = float(line.strip())
                    slot_data[idx].append(value)

        # 计算每个时隙的平均值
        slot_averages = [sum(data) if data else 0 for data in slot_data]

        return slot_averages

    def smooth_data(self, data, window_size=10):
        # 使用移动平均对数据进行平滑处理
        smoothed_data = []
        for i in range(len(data)):
            start_index = max(0, i - window_size + 1)
            end_index = i + 1
            window_values = data[start_index:end_index]
            smoothed_value = sum(window_values) / len(window_values) if window_values else 0
            smoothed_data.append(smoothed_value)

        return smoothed_data

    def write_smoothed_data(self, output_filename, window_size=10):
        # 获取每个时隙的平均值
        slot_averages = self.calculate_slot_averages()

        # 对数据进行平滑处理
        smoothed_data = self.smooth_data(slot_averages, window_size=window_size)

        # 将平滑后的结果写入文件
        with open(output_filename, 'w') as file:
            for value in smoothed_data:
                file.write(f"{value}\n")

        print(f"平滑后的数据已成功写入文件: {output_filename}")

    def plot_total_throughput(self, window_size=2000):
        # 获取每个时隙的平均值
        slot_averages = self.calculate_slot_averages()

        # 对数据进行平滑处理
        smoothed_data = self.smooth_data(slot_averages, window_size=window_size)

        # 计算总吞吐量
        max_iter = len(slot_averages)
        total_throughput = np.cumsum(smoothed_data) / np.arange(1, max_iter + 1)

        # 绘制总吞吐量随时间变化的曲线图
        plt.figure(figsize=(10, 6))
        plt.plot(np.arange(1, max_iter + 1), total_throughput, color='#236B8E', lw=1.2, label='Total Throughput')
        plt.xlabel('Time Step')
        plt.ylabel('Total Throughput')
        plt.title('Total Throughput over Time')
        plt.legend()
        plt.grid(True)
        plt.show()

# 示例用法

filenames = [
        'rewards/agent_len5e4_M40.txt',
        'rewards/TDMA_len5e4_M40.txt',
        'rewards/SA_len5e4_M40.txt'
]
output_filename = 'rewards/total.txt'

calculator = CalculateAverages(filenames)
calculator.write_smoothed_data(output_filename, window_size=5000)  # 使用窗口大小为2000的移动平均进行平滑操作
calculator.plot_total_throughput(window_size=5000)  # 绘制总吞吐量随时间变化的曲线图